package com.amumu.drama.common.config.redis;

import com.google.common.base.Charsets;
import com.google.common.hash.Hashing;

/**
 * 基于令牌桶算法的限流器
 * <p>
 * 例如针对指定ip限流5分钟内请求10次：
 * <p>
 * JedisRateLimiterConfig config=new JedisRateLimiterConfig("ipratelimit",10,5*60);
 * if(jedisTemplate.rateLimit(config).tryAcquire(ip)){
 * // 获得请求令牌
 * }
 * <p>
 * <p>
 * see https://gist.github.com/ptarjan/e38f45f2dfe601419ca3af937fff574d#request-rate-limiter
 * see https://stripe.com/blog/rate-limiters
 * <p>
 * Created by freethinker on 2020-07-01.
 */
public class JedisRateLimiter {

    // @formatter:off
    private static final String SCRIPT=
        "local tokens_key = KEYS[1] " +             // 限流名
        "local timestamp_key = KEYS[2] " +          // 限流名的时间戳
        "" +
        "local rate = tonumber(ARGV[1]) " +         // 每秒钟生成令牌数
        "local capacity = tonumber(ARGV[2]) " +     // 桶容量
        "local now = tonumber(ARGV[3]) " +          // 当前时间，单位秒
        "local requested = tonumber(ARGV[4]) " +    // 当前请求数
        "" +
        "local fill_time = capacity/rate " +        // 填充时间
        "local ttl = math.floor(fill_time*2) " +    // 填充时间*2为key的ttl
        "" +
        "local last_tokens = tonumber(redis.call('get', tokens_key)) " +
        "if last_tokens == nil then " +
        "  last_tokens = capacity " +              // 如果取不到，则容量为上次可用令牌数
        "end " +
        "" +
        "local last_refreshed = tonumber(redis.call('get', timestamp_key)) " +
        "if last_refreshed == nil then " +
        "  last_refreshed = 0 " +                  // 如果取不到，则上次时间为0
        "end " +
        "" +
        "local delta = math.max(0, now-last_refreshed) " +   // 上次的时间差
        "local filled_tokens = math.min(capacity, last_tokens+(delta*rate)) " +  // 填充后可用令牌数
        "local allowed = filled_tokens >= requested " +      // 如果可用令牌数大于请求量，则允许请求
        "local new_tokens = filled_tokens " +                // 最新的领取可用数
        "if allowed then " +
        "  new_tokens = filled_tokens - requested " +        // 如果允许请求，则扣减令牌数
        "end " +
        "" +
        "redis.call('setex', tokens_key, ttl, new_tokens) " +
        "redis.call('setex', timestamp_key, ttl, now) " +
        "" +
        "return allowed";                                    // 是否允许请求，1:允许请求，nul:不允许
    // @formatter:on

    /** 限流器实例名,具体名字的md5取中间8位 */
    private String limitRateName;

    private JedisTemplate jedisTemplate;

    /** 令牌桶补充速率，即允许的消耗速度 */
    private double replenishRate;

    /** 桶的容量，即个数 */
    private int burstCapacity;

    JedisRateLimiter(JedisRateLimiterConfig config, JedisTemplate jedisTemplate) {
        this.limitRateName = Hashing.md5().hashString(config.getLimitRateName(), Charsets.UTF_8).toString().substring(8, 16);
        this.jedisTemplate = jedisTemplate;
        this.replenishRate = 1.0 * config.getLimitCount() / config.getLimitSpanSeconds();
        this.burstCapacity = config.getLimitCount();
    }

    /**
     * 获取请求令牌
     *
     * @param limitResource 请求资源
     * @return true:获取令牌，false:未获得令牌
     */
    public boolean tryAcquire(Object limitResource) {
        return tryAcquire(limitResource, 1);
    }

    /**
     * 获取请求令牌
     *
     * @param limitResource 请求资源
     * @param permits       一次性获取的令牌数
     * @return true:获取令牌，false:未获得令牌
     */
    public boolean tryAcquire(Object limitResource, int permits) {
        if (permits <= 0) {
            return true;
        }
        String resource = Hashing.md5().hashString(String.valueOf(limitResource), Charsets.UTF_8).toString().substring(8, 16);
        return jedisTemplate.execute(jedis -> {
            Object eval = jedis
                    .eval(SCRIPT, 2, "rl_" + limitRateName + "_" + resource, "rlts_" + limitRateName + "_" + resource, "" + replenishRate, "" + burstCapacity, "" + (System.currentTimeMillis() / 1000),
                            "" + permits);
            return null != eval;
        });
    }
}
